import torch
from algorithms.mat.algorithm.ma_transformer import MultiAgentTransformer


# MAT策略类。包装actor和critic网络来计算动作和价值函数预测。
class TransformerPolicy:

    def __init__(self, args, obs_shape, act_shape, num_agents, device=torch.device("cpu")):
        """
        初始化Transformer策略网络
        
        Args:
            args: 配置参数
            obs_shape: 观测形状 (36,) - 36维状态向量
            act_shape: 动作数量 7
            num_agents: 智能体数量
            device: 计算设备
        """
        super(TransformerPolicy, self).__init__()

        # 基本参数
        self.obs_shape = obs_shape  # 观测形状 (36,)
        self.act_shape = act_shape  # 动作数量 7
        self.num_agents = num_agents  # 智能体数量
        self.lr = args.lr  # 学习率
        self.weight_decay = args.weight_decay  # 权重衰减
        self.device = device  # 计算设备
        self.tpdv = dict(dtype=torch.float32, device=device)  # tensor参数（数据类型和设备）

        # 使用支持36维状态的MAT模型
        self.transformer = MultiAgentTransformer(
            action_dim=self.act_shape,  # 7维动作空间
            n_agent=num_agents,
            n_block=args.n_block,
            n_embd=args.n_embd,
            n_head=args.n_head,
            device=device,
            obs_dim=36,  # 36维状态向量
            dropout_rate=args.dropout_rate  # Dropout比率
        )

        # 端到端训练的优化器 - 使用AdamW和权重衰减
        self.optimizer = torch.optim.AdamW(
            self.transformer.parameters(),
            lr=self.lr,
            weight_decay=self.weight_decay
        )

        # 学习率调度器
        self.scheduler = torch.optim.lr_scheduler.StepLR(
            self.optimizer,
            step_size=1000,
            gamma=0.9
        )

        # 梯度裁剪参数
        self.max_grad_norm = 0.5

    def get_actions(self, obs, available_actions=None, deterministic=False):
        """
        计算给定输入的动作和价值函数预测 - 端到端训练
        
        Args:
            obs: 形状为 (batch_size, n_agents, 36) 的36维状态
            available_actions: 可用动作掩码
            deterministic: 是否确定性采样
            
        Returns:
            actions: 选择的动作
            action_log_probs: 动作对数概率
            values: 价值函数预测
        """
        # 使用支持36维状态的MAT模型获取动作 - 端到端训练
        actions, action_log_probs, values = self.transformer.get_actions(
            obs, available_actions, deterministic
        )

        return actions, action_log_probs, values

    def evaluate_actions(self, obs, actions, available_actions=None, active_masks=None):
        """
        评估动作的价值和对数概率 - 端到端训练
        
        Args:
            obs: 形状为 (batch_size, n_agents, 36) 的36维状态
            actions: 要评估的动作
            available_actions: 可用动作掩码
            active_masks: 活跃掩码
            
        Returns:
            values: 价值函数预测
            action_log_probs: 动作对数概率
            dist_entropy: 分布熵
        """
        # 使用支持36维状态的MAT模型评估动作 - 端到端训练
        action_log_probs, values, dist_entropy = self.transformer.forward(
            obs, actions, available_actions
        )

        return values, action_log_probs, dist_entropy

    # 保存模型
    def save(self, save_dir, episode):
        torch.save(self.transformer.state_dict(), str(save_dir) + "/transformer_" + str(episode) + ".pt")  # 保存模型参数

    # 恢复模型
    def restore(self, model_dir):
        transformer_state_dict = torch.load(model_dir)  # 加载模型参数
        self.transformer.load_state_dict(transformer_state_dict)  # 恢复模型参数

    # 设置为训练模式
    def train(self):
        self.transformer.train()

    def update(self, loss):
        """
        执行端到端训练更新
        
        Args:
            loss: 总损失
        """
        # 清空梯度
        self.optimizer.zero_grad()

        # 反向传播
        loss.backward()

        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(self.transformer.parameters(), self.max_grad_norm)

        # 参数更新
        self.optimizer.step()

        # 学习率调度
        self.scheduler.step()

    def get_trainable_params(self):
        """获取可训练参数数量"""
        return sum(p.numel() for p in self.transformer.parameters() if p.requires_grad)

    def save_model(self, path):
        """保存模型"""
        torch.save({
            'transformer_state_dict': self.transformer.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
        }, path)

    def load_model(self, path):
        """加载模型"""
        checkpoint = torch.load(path, map_location=self.device)
        self.transformer.load_state_dict(checkpoint['transformer_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

    # 设置为训练模式
    def train(self):
        self.transformer.train()

    # 设置为评估模式
    def eval(self):
        self.transformer.eval()
